##PBAR_ZspecMC.py

import csv
import numpy as np
import numpy.matlib
import datetime
import time

# BASIC STUFF

def 

def ReadZspec(filenameList):
    dat = list()
    for ii in range(len(fullfilenameList)):
        temp = np.genfromtxt(fullfilenameList[ii], \
                                 delimiter=',', \
                                 skip_header = 0, \
                                 skip_footer = 0, \
                                 dtype = 'uint32')
        if (temp.shape[0] == 257):  # shave off the table labels if present
            temp = temp[1:,:]
        if (temp.shape[1] == 138):
            temp = temp[:,1:]
        dat.append(temp)
    return dat

def ReadRad(basepath):
    # read in the radiography data
    radMap = np.genfromtxt(basepath + '\\' + 'arrayResponse.txt', delimiter = '\t', skip_header = 1, dtype = str)

    # load all the radiography data fount in radMap and put into a dictionary
    datRad = dict()
    for ii in range(0, radMap.shape[0]):
        dsName = radMap[ii,0]
        if (radMap[ii,3] != ''):
            radName = radMap[ii,3]
            fullFileName = basepath + '\\' + 'Array-Response' + '\\' + radName + '.csv'
            temp = np.genfromtxt(fullFileName, delimiter = '\t', skip_header = 1, dtype = str)
            datRad[dsName] = temp.astype(float)
    return datRad, radMap

def GetDatasetInformation(infoFilename, filenameList):
    fid = open(infoFilename, 'r');
    csvReaderObj = csv.reader(fid, delimiter = '\t')
    datasetName = []
    acquisitionTime = []
    timeStamp = []
    timeNum = []
    header = []

    for lineIn in csvReaderObj:
        datasetName.append(lineIn[0])
        timeStamp.append(lineIn[1])
        acquisitionTime.append(lineIn[2])
        timeNum.append(lineIn[3])
        header.append(lineIn[4])

    acquisitionTime = np.array(acquisitionTime)
    acquisitionTime = acquisitionTime.astype(float)
    datasetName = np.array(datasetName)
    timeStamp = np.array(timeStamp)
    timeNum = np.array(timeNum)
    timeNum = timeNum.astype(float)
    header = np.array(header)

    # Get the information about the dataset
    datasetDescription = [] # header
    datasetTimeStr = [] # dataset time stamp string
    datasetTime = [] # dataset time stamp string
    datasetTimeNum = [] # epoch start time
    datasetAcquisitionTime = [] # acquisition  time

    for ii in range(0, len(filenameList)):
        index = np.where(datasetName == filenameList[ii])[0][0]
        datasetDescription.append(header[index])
        datasetTimeStr.append(timeStamp[index])
        datasetTimeNum.append(timeNum[index])
        datasetTime.append(time.localtime(timeNum[index]))
        datasetAcquisitionTime.append(acquisitionTime[index])

    datasetDescription = np.array(datasetDescription)
    datasetTimeStr = np.array(datasetTimeStr)
    datasetTimeNum = np.array(datasetTimeNum)
    datasetAcquisitionTime = np.array(datasetAcquisitionTime)

    return (datasetDescription, datasetAcquisitionTime, datasetTime, datasetTimeNum, datasetTimeStr)


def CreateGainCorrectionFile(filenameList):
    # read in all the dataset files
    
    # read in the information file

    # calculate the gain correction parameters

    # calculate the 
    1+1




def SimulateAll(filenameList, binList, timeList, statsList, numTrials, outputDir):
    for filename in range(len(filenameList)):
        
        # read in file
        dat = ReadZspec(filename)

        # get dataset information from header csv file
        (descript, acqTime, datasetTime, datasetTimeNum, datasetTimeStr) = \
                             GetDatasetInformation(infoFilename, filenameList)
        # read in gain correction file

        # cycle throught list of data set acquisition times to simulate
        for t in timeList:
            1+1
            # apply gain correction on matrix
            
            # generate the spectra matrix
            
            # save the spectrum data to file <dataset name><detector index><timing information><current time stamp>.txt
            # cycle through the list of energy bins
            for bins in binList:
                1+1
                # rebin the data


# need the datasetSummary file? 
def CalculateStatsAll(filenameList, statsList):
    1+1
    # cycle through datasets
        # load all files for a dataset
            #  cycle through different binning structure
                # rebin the spectra
                # calculate the stat values
                # save to file
    
def CalculateStats(spectrumMat, spectrumBins, statsList, gainCorrectionMat):
    # initialize the array that holds the stats
    statsMC = dict()

    statsListDefault = ('binSum', \
                 'binMean', 'binSTD', 'binSTD_binMean',\
                'q_10', 'q_25', 'q_50', 'q_75', 'q_90', \
                'q_range', 'q_range_ratio', 'binMeanSq', 'binMean_q50', \
                 \
                'transmission')

    for i in range(0,len(statsListDefault)):
        statsMC[statsListDefault[i]] = np.zeros(spectrumMat.shape[1])

    binArray = np.matlib.repmat(spectrumBins, spectrumMat.shape[1],1).T

    # CALCULATE THE STATISTICS

    # remove the lower bins by setting it to zero
    datMC[0:lowerBinThreshold,:] = 0

    # Find the quantile bins
    datTempCummulativeSum = datMC.cumsum(axis = 0) # cumulative sum
    datTempCummulativeSum = datTempCummulativeSum / np.matlib.repmat(datTempCummulativeSum[-1,:], 256, 1) # normalize the cumulative sum
    
    statsMC['q_10'][ii,:] = np.argmin(abs(datTempCummulativeSum - 0.10), axis = 0)
    statsMC['q_25'][ii,:] = np.argmin(abs(datTempCummulativeSum - 0.25), axis = 0)
    statsMC['q_50'][ii,:] = np.argmin(abs(datTempCummulativeSum - 0.50), axis = 0)
    statsMC['q_75'][ii,:] = np.argmin(abs(datTempCummulativeSum - 0.75), axis = 0)
    statsMC['q_90'][ii,:] = np.argmin(abs(datTempCummulativeSum - 0.90), axis = 0)
    # quantile range
    statsMC['q_range'][ii,:] = statsMC['q_75'][ii,:] - statsMC['q_25'][ii,:]
    statsMC['q_range_ratio'][ii,:] = statsMC['q_range'][ii,:] / statsMC['q_50'][ii,:]

    # calculate the mean and spread    
    statsMC['binMean'][ii,:] = (binArray * datMC).sum(axis = 0).astype(float) / datMC.sum(axis = 0).astype(float)
    statsMC['binMeanSq'][ii,:] = ((binArray**2) * datMC).sum(axis = 0).astype(float) / datMC.sum(axis = 0).astype(float)
    
    statsMC['binSum'][ii,:] = datMC.sum(axis = 0).astype(float)

    statsMC['binMean_q50'][ii,:] = statsMC['binMean'][ii,:] - statsMC['q_50'][ii,:]

    temp = np.matlib.repmat(statsMC['binMean'][ii,:], datMC.shape[0],1)  # array of means, 256 x number detectors

    statsMC['binSTD'][ii,:] = np.sqrt(  ( ((binArray - temp)**2) * datMC).sum(axis = 0).astype(float) / datMC.sum(axis = 0).astype(float)  )

    statsMC['binSTD_binMean'][ii,:] = statsMC['binSTD'][ii,:] / statsMC['binMean'][ii,:]
    
    # MULTIBIN PARAMETERS
    # Calculate histogram splits - no rebinning so no correction
    statsMC['multibin_0_10'][ii,:] = datMC[0:10,:].sum(axis=0)  # bins 0 to 9, so first ten bins
    statsMC['multibin_10_256'][ii,:] = datMC[10:,:].sum(axis=0)
    statsMC['multibin_0_20'][ii,:] = datMC[0:20,:].sum(axis=0)
    statsMC['multibin_20_256'][ii,:] = datMC[20:,:].sum(axis=0)

    statsMC['multibin_10_ratio'][ii,:]  = statsMC['multibin_0_10'][ii,:]  /statsMC['multibin_10_256'][ii,:]
    statsMC['multibin_20_ratio'][ii,:]  = statsMC['multibin_0_20'][ii,:]  /statsMC['multibin_20_256'][ii,:]

    # transmission value - sum spectra / acquisition time
    statsMC['transmission'][ii,:] = datMC.sum(axis = 0) / acquisitionTimeMC
    return statsMC

def SimulateSpectrum(spectrum, binsInput, binsOutput, numberTrials, t0, t):
    """asdf"""
    # container for the MC spectra
    datMC = np.zeros((len(binsOutput), numberTrials))

    # calculate the total number of counts in spectrum
    totalCounts = spectrum.sum().astype(float)
    
    # calculate mean number of counts expected in simulated data
    numberCountsMC = totalCounts * t / t0  # assume same pulse rate

    # create pdf
    temp = np.interp(binsOutput, binsInput, spectrum.astype(float))
    probdist = temp / temp.sum()

    binEdges = np.arange(257)
    # Run Trials
    for trialNo in range(numberTrials):
        # randomize number of counts
        countsTemp = round(numberCountsMC + np.random.randn() * np.sqrt(numberCountsMC))
        # Calculate bins of x-ray
        bins = np.random.choice(binsOutput, size=countsTemp, replace=True, p=probdist)
        # Make histogram of bins
        datMC[:,trialNo], binEdges = np.histogram(bins, bins = binEdges)

    return datMC
